import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter


class MySequential(nn.Sequential):
    def __init__(self):
        super(MySequential, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 32, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        return self.model1(x)


class MyNet1(nn.Module):
    def __init__(self):
        super(MyNet1, self).__init__()
        self.conv2d1 = nn.Conv2d(3, 32, kernel_size=5, padding=2)
        print("Weights:", self.conv2d1.weight.shape)
        print("Bias:", self.conv2d1.bias.shape)
        # print("Bias:", self.conv2d1.bias)
        self.maxpool2d1 = nn.MaxPool2d(kernel_size=2)
        self.conv2d2 = nn.Conv2d(32, 32, kernel_size=5, padding=2)
        self.maxpool2d2 = nn.MaxPool2d(kernel_size=2)
        self.conv2d3 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
        self.maxpool2d3 = nn.MaxPool2d(kernel_size=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, input_):
        print('input before conv2d1', input_.shape)
        input_ = self.conv2d1(input_)
        print('input after conv2d1', input_.shape)
        input_ = self.maxpool2d1(input_)
        print('input after maxpool2d1', input_.shape)
        input_ = self.conv2d2(input_)
        print('input after conv2d2', input_.shape)
        input_ = self.maxpool2d2(input_)
        print('input after maxpool2d2', input_.shape)
        input_ = self.conv2d3(input_)
        print('input after conv2d3', input_.shape)
        input_ = self.maxpool2d3(input_)
        print('input after maxpool2d3', input_.shape)
        input_ = self.flatten(input_)
        print('input after flatten', input_.shape)
        input_ = self.fc1(input_)
        print('input after fc1', input_.shape)
        input_ = self.fc2(input_)
        print('input after fc2', input_.shape)
        return input_


net = MyNet1()
seq = MySequential()
x = torch.zeros((64, 3, 32, 32))
# y = net(x)
y = seq(x)
writer = SummaryWriter("../logs")
writer.add_graph(net, x)
print('-' * 50)
print(y.shape)
print(y)
